package com.google.code.sfa.nextgen;

import java.util.*;
import java.io.*;
import org.biojava.bio.symbol.*;
import org.biojava.bio.seq.DNATools;
import com.google.code.sfa.util.*;
import com.google.code.sfa.parse.*;
import com.google.code.sfa.tigr.*;


public class NextGenSequencer {


    public void assemble( List<Trace> traces, SymbolList reference, SymbolList subject) {

	Indexes indexes = new Indexes();

	for( Trace t : traces) {
	    //System.out.println( t.toString() );
	    indexes.add( t);
	}

	AcceptableMatches goodMatches = new AcceptableMatches();

	int subListLowerBound = 1, subListUpperBound;
	for( int i = 1; i+32 < reference.length(); i++) {
	    if( i > 3) {
		subListLowerBound++;
	    }
	    OverlapAlignment aligner = new OverlapAlignment();
	    MatchCriteria mc = new MatchCriteria(32, 4, 0.8875, 2);
	    Set<Trace> matches = indexes.findMatches( reference, i).keySet();
	    for( Trace match : matches) {
		
		if( match.getSequence().length() > reference.length() - i + 1) {
		    continue;
		}

		if( i < 3) {
		    subListUpperBound = subListLowerBound + match.getSequence().length() + i;
		} else {
		    subListUpperBound = subListLowerBound + match.getSequence().length() + 3;
		}
		if( subListUpperBound > reference.length() ) {
		    subListUpperBound = reference.length();
		}
		
		aligner.align( match.getSequence(), reference.subList(subListLowerBound, subListUpperBound));

		if( mc.testMatch( aligner)) {
		    Match m = new Match();
		    //System.out.println( aligner.offset);
		    m.refPos = aligner.offset + subListLowerBound - 1;
		    m.errors = mc.errors;
		    m.alignment = aligner;
		    goodMatches.add( match, m);
		    aligner = new OverlapAlignment();
		}
	    }
	} 
	
	Set<Trace> tracesThatMatch = goodMatches.keySet();

	Trace referenceTrace = new Trace("REFERENCE", "REFERENCE", reference, true);

	PositionInserts inserts = new PositionInserts();
	for( Trace t : tracesThatMatch) {
	    setupTrace( t, goodMatches.bestMatch(t), inserts);
	}

	Match m;
	for( Trace t : tracesThatMatch) {
	    m = goodMatches.bestMatch(t);
	    m.refPos = insertGapsInTrace( t, m.refPos, inserts);
	}
	insertGapsInTrace( referenceTrace, 1, inserts);

	Contig c = new Contig( referenceTrace);
	for( Trace t : tracesThatMatch) {
	    c.set( t, goodMatches.bestMatch(t).refPos);
	}
	System.out.println( c.toString());

	OverlapAlignment finalAlignment = new OverlapAlignment();
	finalAlignment.align( c, subject);
	SymbolList querySeq = finalAlignment.alignment.symbolListForLabel("query");
	SymbolList subjectSeq = finalAlignment.alignment.symbolListForLabel("subject");
	System.out.println( "REAL " + subjectSeq.seqString());
	System.out.print("EXP  ");
	for( int i = 1; i <= querySeq.length(); i++) {
	    try {
		if( querySeq.symbolAt(i) == subjectSeq.symbolAt(i)) {
		    System.out.print(".");
		} else {
		    System.out.print( DNATools.dnaToken( querySeq.symbolAt(i)));
		}
	    } catch( IllegalSymbolException e) {
		System.out.print("?");
	    }
	}
	System.out.println();

	/*
	PriorityQueue<Match> matchQueue = new PriorityQueue<Match>();
	Match current;
	while( (current = matchQueue.poll() ) != null) {
	    System.out.println( current.refPos);	    
	}
	*/
    }

    private int insertGapsInTrace( Trace t, int first, PositionInserts inserts) {
	int gapsToAdd, gapsAdded;
	int referencePos = first + 1; // don't care about gaps added before the first position
	int seqPos = 2;
	while( seqPos < t.getSequence().length()) {
	    gapsToAdd = inserts.gapsToAdd( t, referencePos);
	    gapsAdded = inserts.gapsAdded( t, referencePos);
	    while( gapsToAdd > 0) {
		t.insertGap( seqPos);
		gapsToAdd--;
		seqPos++;
	    }
	    seqPos += gapsAdded + 1;
	    referencePos++;
	}
	int newFirst = first;
	for( int i = 1; i <= first; i++) {
	    newFirst += inserts.gapsToAdd( t, i);
	}
	return newFirst;
    }

    private void setupTrace( Trace t, Match m, PositionInserts inserts) {
	Alignment alignment = m.alignment.alignment;
	Symbol gap = DNATools.getDNA().getGapSymbol();
	int existingPos = m.refPos;

	for( int i = 1; i < alignment.length(); i++) {
	    if( alignment.symbolAt("subject", i) == gap) {
		inserts.add( t, existingPos);
	    } else {
		existingPos++;
	    }

	    if( alignment.symbolAt("query",i) == gap) {
		t.insertGap(i);
	    }
	}

    }

}

class PositionInserts {
    
    private Map<Integer,PositionInsert> inserts = new HashMap<Integer, PositionInsert>();

    public void add( Trace t, int pos) {
	PositionInsert pi;
	if( inserts.containsKey( pos)) {
	    pi = inserts.get(pos);
	} else {
	    pi = new PositionInsert();
	    inserts.put( pos, pi);
	}
	pi.add(t);
    }

    public int gapsToAdd( Trace t, int pos) {
	if( inserts.containsKey( pos) ) {
	    return inserts.get(pos).gapsToAdd(t);
	} else {
	    return 0;
	}
    }
    public int gapsAdded( Trace t, int pos) {
	if( inserts.containsKey( pos) ) {
	    return inserts.get(pos).gapsAdded(t);
	} else {
	    return 0;
	}
    }	

}

class PositionInsert {
    private Map<Trace,Integer> inserts = new HashMap<Trace,Integer>();
    private int maxGaps = 0;

    public void add( Trace t) {
	int gaps;
	if( inserts.containsKey(t) ) {
	    gaps = inserts.get(t) + 1;
	} else {
	    gaps = 1;
	}
	inserts.put(t, gaps);
	if( maxGaps < gaps) {
	    maxGaps = gaps;
	}
    }

    public int gapsToAdd( Trace t) {
	if( inserts.containsKey(t) ) {
	    return maxGaps - inserts.get(t);
	} else {
	    return maxGaps;
	}
    }
    public int gapsAdded( Trace t) {
	if( inserts.containsKey(t) ) {
	    return inserts.get(t);
	} else {
	    return 0;
	}
    }
}


class Match implements Comparable<Match> {
    public int refPos;
    public int errors;
    public OverlapAlignment alignment;

    public int compareTo( Match other) {
	return other.refPos - refPos;
    }
}

class AcceptableMatches {
    private Map<Trace,Match> allMatches = new HashMap<Trace,Match>();

    public void add( Trace t, Match m) {
	List<Match> matches;
	if( !allMatches.containsKey( t) ) {
	    allMatches.put( t, m);
	} else {
	    Match existingMatch = allMatches.get( t);
	    if( existingMatch.errors < m.errors) {
		allMatches.put( t, m);
	    }
	}
    }

    public Set<Trace> keySet() {
	return allMatches.keySet();
    }

    public Match bestMatch(Trace t) {
	List<Match> matches;
	if( !allMatches.containsKey( t) ) {
	    return null;
	} else {
	    return allMatches.get( t);
	}
    }
}
	    
	    
	